#!/usr/bin/env python3

import argparse
import filecmp
import os
import shutil
import sys

BLUE = '\033[34m'
GREEN = '\033[32m'
RED = '\033[31m'


def color_text(color, t):
    if os.isatty(1):
        return color + t + '\033[0m'
    else:
        return t


def print_color(color, t):
    print(color_text(color, t))


def are_same_files(a, b):
    return filecmp.cmp(a, b, shallow=True)


def copy_file(src, dst, overwrite):
    if os.path.exists(dst):
        if are_same_files(src, dst):
            print_color(BLUE, os.path.basename(dst) + ' already exists. No need to copy it.')
            return
        elif not overwrite:
            print_color(RED, dst + ' already exists, but it differs')
            print_color(RED, '  from ' + src + ' .')
            print_color(RED, '  ' + os.path.basename(dst) + ' will not be copied.')
            print_color(RED,'  Use the -o/--overwrite option to copy this file anyway.')
            return

    try:
        shutil.copyfile(src, dst)
        print_color(GREEN, 'Copying ' + os.path.basename(dst))
    except:
        print_color(RED,'Error: Could not copy ' + src + ' to ' + dst)


def link_file(src, dst, overwrite):
    if os.path.exists(dst):
        if are_same_files(src, dst):
            print_color(BLUE, os.path.basename(dst) + ' already exists. No need to link it.')
            return
        elif not overwrite:
            print_color(RED, dst + ' already exists, but it differs')
            print_color(RED, '  from ' + src + ' .')
            print_color(RED, '  ' + os.path.basename(dst) + ' will not be linked.')
            print_color(RED, '  Use the -o/--overwrite option to link this file anyway.')
            return
        else:
            os.remove(dst)

    try:
        os.symlink(src, dst)
        print_color(GREEN, 'Linking ' + os.path.basename(dst))
    except:
        print_color(RED, 'Error: Could not link ' + src +' to '+ dst)


def missing_files(mpas_root, required_files):
    file_list = []
    for f in required_files:
        if not os.path.isfile(os.path.join(mpas_root, f)):
            file_list.append(f)

    if len(file_list) == 0:
        return None
    else:
        return file_list


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('run_dir',
                        help='the run directory to set up')
    parser.add_argument('-o', '--overwrite',
                        action='store_true',
                        help='overwrite any existing files and symbolic links')
    core_group = parser.add_mutually_exclusive_group()
    core_group.add_argument('-i', '--init-only',
                            action='store_true',
                            help='only set up files for the init_atmosphere core')
    core_group.add_argument('-a', '--atm-only',
                            action='store_true',
                            help='only set up files for the atmosphere core')
    args = parser.parse_args()

    mpas_root = os.path.abspath(os.path.join(os.path.dirname(sys.argv[0]), '..', '..'))
    run_dir = os.path.abspath(args.run_dir)
    print('Root MPAS-Model directory: ' + mpas_root)
    print('Run directory to set up: ' + run_dir)
    if args.init_only:
        print('Only setting up files for the init_atmosphere core...')
    elif args.atm_only:
        print('Only setting up files for the atmosphere core...')
    else:
        print('Setting up files for both the init_atmosphere and atmosphere cores...')

    print('---')

    # Lookup tables used by WRF physics schemes
    physics_wrf = ['src/core_atmosphere/physics/physics_wrf/files/' + f
        for f in [
            'CAM_ABS_DATA.DBL',
            'CAM_AEROPT_DATA.DBL',
            'CCN_ACTIVATE_DATA',
            'GENPARM.TBL',
            'LANDUSE.TBL',
            'OZONE_DAT.TBL',
            'OZONE_LAT.TBL',
            'OZONE_PLEV.TBL',
            'RRTMG_LW_DATA',
            'RRTMG_LW_DATA.DBL',
            'RRTMG_SW_DATA',
            'RRTMG_SW_DATA.DBL',
            'SOILPARM.TBL',
            'VEGPARM.TBL',
            ]
        ]

    # Lookup tables used by Noah-MP
    physics_noahmp = ['src/core_atmosphere/physics/physics_noahmp/parameters/' + f
        for f in [
            'NoahmpTable.TBL',
            ]
        ]

    # Optional lookup tables that are generated by the build_tables utility
    opt_physics_wrf = ['src/core_atmosphere/physics/physics_wrf/files/' + f
        for f in [
            'MP_THOMPSON_freezeH2O_DATA.DBL',
            'MP_THOMPSON_QIautQS_DATA.DBL',
            'MP_THOMPSON_QRacrQG_DATA.DBL',
            'MP_THOMPSON_QRacrQS_DATA.DBL',
            ]
        ]

    init_files_to_link = ['init_atmosphere_model']
    init_files_to_copy = ['default_inputs/' + f
        for f in [
            'namelist.init_atmosphere',
            'streams.init_atmosphere',
            ]
        ]

    atm_files_to_link = ['atmosphere_model', 'build_tables']
    atm_files_to_link = atm_files_to_link + physics_wrf + physics_noahmp
    atm_opt_files_to_link = opt_physics_wrf
    atm_files_to_copy = ['default_inputs/' + f
        for f in [
            'namelist.atmosphere',
            'streams.atmosphere',
            'stream_list.atmosphere.output',
            'stream_list.atmosphere.diagnostics',
            'stream_list.atmosphere.surface',
            'stream_list.atmosphere.diag_ugwp',
            ]
        ]

    # Check if the init_atmosphere core has been compiled in the MPAS-Model directory
    if not args.atm_only:
        files = missing_files(mpas_root, init_files_to_link + init_files_to_copy)
        if files:
            print_color(RED, 'Error: The init_atmosphere core does not appear to have been successfully built')
            print_color(RED, '   in ' + mpas_root + '. The following files were expected there,')
            print_color(RED, '   but were not found:')
            for f in files:
                print_color(RED, '           ' + f)
            sys.exit(1)

    # Check if the atmosphere core has been compiled in the MPAS-Model directory
    if not args.init_only:
        files = missing_files(mpas_root, atm_files_to_link + atm_files_to_copy)
        if files:
            print_color(RED, 'Error: The atmosphere core does not appear to have been successfully built')
            print_color(RED, '   in ' + mpas_root + ' .')
            print_color(RED, '   The following files were expected there but were not found:')
            for f in files:
                print_color(RED, '           ' + f)
            sys.exit(1)

    try:
        os.makedirs(run_dir)
        print_color(GREEN, 'Creating directory ' + run_dir)
    except FileExistsError:
        print_color(BLUE, 'Directory ' + run_dir + ' already exist. No need to create it.')
    except:
        print_color(RED, 'Error: Could not create directory ' + run_dir)
        sys.exit(1)

    if not os.access(run_dir, os.W_OK):
        print_color(RED, 'Error: No write permission for directory ' + run_dir)
        sys.exit(1)

    if not args.atm_only:
        print('')
        print_color(GREEN, 'Setting up init_atmosphere core files...')
        print_color(GREEN, '----------------------------------------')
        for f in init_files_to_copy:
            copy_file(os.path.join(mpas_root, f),
                      os.path.join(run_dir, os.path.basename(f)),
                      args.overwrite)

        for f in init_files_to_link:
            link_file(os.path.join(mpas_root, f),
                      os.path.join(run_dir, os.path.basename(f)),
                      args.overwrite)

    if not args.init_only:
        print('')
        print_color(GREEN, 'Setting up atmosphere core files...')
        print_color(GREEN, '----------------------------------------')
        for f in atm_files_to_copy:
            copy_file(os.path.join(mpas_root, f),
                      os.path.join(run_dir, os.path.basename(f)),
                      args.overwrite)

        for f in atm_files_to_link:
            link_file(os.path.join(mpas_root, f),
                      os.path.join(run_dir, os.path.basename(f)),
                      args.overwrite)

        for f in atm_opt_files_to_link:
            if os.path.isfile(mpas_root + '/' + f):
                link_file(os.path.join(mpas_root, f),
                          os.path.join(run_dir, os.path.basename(f)),
                          args.overwrite)
